package org.totoro.structure.tree.avl;

import org.totoro.structure.tree.Tree;

import java.util.ArrayList;
import java.util.Objects;

/**
 * 平衡二叉树
 *
 * @author daocr
 * @date 2020/4/12
 */
public class AvlTree<E extends Comparable<E>> implements Tree<E> {

    private int size;

    private TreeNode<E> root;

    public AvlTree() {

    }

    /**
     * 初始化
     *
     * @param element
     */
    public AvlTree(E element) {
        root = new TreeNode<>(element);
        size++;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean contains(E element) {
        return false;
    }

    private TreeNode<E> find(TreeNode<E> node, E element) {

        if (Objects.isNull(node)) {
            return null;
        }

        if (node.element.compareTo(element) == 0) {
            return node;
        }

        TreeNode<E> findLeft = find(node.left, element);

        if (Objects.nonNull(findLeft)) {
            return findLeft;
        }

        return find(node.right, element);

    }


    @Override
    public boolean add(E element) {

        if (Objects.isNull(root)) {
            root = new TreeNode<>(element);
            size++;
            return true;
        }
        add(root, element);
        size++;
        return false;
    }

    /**
     * 添加
     *
     * @param node
     * @param element
     * @return
     */
    private TreeNode<E> add(TreeNode<E> node, E element) {

        if (Objects.isNull(node)) {
            return new TreeNode<>(element);
        }

        E curNodeElement = node.element;

        // 小于等于当前节点
        if (element.compareTo(curNodeElement) <= 0) {
            // 自平衡
            node.left = add(node.left, element);

            return balance(node);
        } else {
            // 大于当前节点就往右边存储
            // 自平衡
            node.right = add(node.right, element);

            return balance(node);
        }

    }

    /**
     * 平衡维护
     * <p>
     * 计算平衡因子
     *
     * @param node
     * @return
     * @see this##calculateFactor(TreeNode)
     */
    private TreeNode<E> balance(TreeNode<E> node) {

        node.height = calculateHeight(node);

        // 计算高度因子
        int calculateFactor = calculateFactor(node);

        /**
         * LL
         *
         * 插入数值 [5] ，导致节点[10]平衡因子大于1，并且左子节点 [8]的平衡因子大于等于0
         *
         *         		 10
         *         		/
         *             8
         *            /
         *           5
         *
         *
         *          * 计算平衡因子
         *          * @see this##calculateFactor(TreeNode)
         *
         */
        if (calculateFactor > 1 && calculateFactor(node.left) >= 0) {
            System.out.println("LL : " + node.element);
            return rightRotate(node);
        }

        //			 (b)
        //			/   \
        //	 	   /     \
        //	 	 (a)     (c)
        //   	/	\	/   \
        //     /	 \ /     \
        //    1      2 3     4


        /**
         *
         *RR
         *
         * 插入 [21] , 导致节点[10]平衡因子小于等于-1，并且左子节点 [8]的平衡因子大于等于[0]
         *
         *    10
         *      \
         *       20
         *        \
         *         21
         *
         *          * 计算平衡因子
         *          * @see this##calculateFactor(TreeNode)
         *
         */
        if (calculateFactor < -1 && calculateFactor(node.right) <= 0) {
            System.out.println("rr : " + node.element);
            return leftRotate(node);
        }

        /**
         *
         * LR
         *
         * 插入数值[6],，导致节点[10]平衡因子大于1，并且左子节点 [8]的平衡因子小于0 ，先执行LL 操作，然后进RR 操作
         *
         *
         *               10
         *              /
         *            8
         *             \
         *              6
         *
         *
         *
         *          计算平衡因子
         *          * @see this##calculateFactor(TreeNode)
         */

        // LR
        if (calculateFactor > 1 && calculateFactor(node.left) < 0) {
            System.out.println("lr: " + node.element);
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }

        /**
         *
         * RL
         *
         * 插入数值[11],，导致节点[10]平衡因子大于1，并且右子节点 [12]的平衡因子大于0 ，先执行RR 操作，然后进LL 操作
         *
         *
         *
         *            10
         *              \
         *              12
         *             /
         *           11
         *
         *
         *
         *          计算平衡因子
         *          * @see this##calculateFactor(TreeNode)
         */

        if (calculateFactor < -1 && calculateFactor(node.right) > 0) {
            System.out.println("rl: " + node.element);
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }


        return node;
    }

    /**
     * 右旋 也叫 rr
     * <p>
     * 数据插入到了 当前节点的右节点的右节点
     *
     * @param node
     * @return
     */
    private TreeNode<E> rightRotate(TreeNode<E> node) {

        //左子树
        TreeNode<E> leftChildNode = node.left;
        // 左子树的右节点
        TreeNode<E> x = leftChildNode.right;

        // 进行向右旋转
        leftChildNode.right = node;
        node.left = x;

        // 更新高度
        node.height = calculateHeight(node);
        leftChildNode.height = calculateHeight(leftChildNode);


        return leftChildNode;

    }


    /**
     * 左旋 也叫 ll
     * <p>
     * 数据插入到了 当前节点的左节点的左节点
     *
     * @param node
     * @return
     */
    private TreeNode<E> leftRotate(TreeNode<E> node) {

        TreeNode<E> rightChildNode = node.right;
        TreeNode<E> x = rightChildNode.left;

        // 向左旋转
        node.right = x;
        rightChildNode.left = node;

        // 计算高度
        node.height = calculateHeight(node);
        rightChildNode.height = calculateHeight(rightChildNode);

        return rightChildNode;

    }

    /**
     * 计算高度
     * <p>
     * 获取左右节点高度，然后取最大值 加 1 返回
     *
     * @param node
     * @return
     */
    private int calculateHeight(TreeNode<E> node) {

        // 获取最大值 然后加一
        return Math.max(getHeight(node.left), getHeight(node.right)) + 1;
    }

    /**
     * 计算高度因子
     * <p>
     * 计算规则：左节点高度减去右节点高度，然后获取绝对值返回
     *
     * @param node
     * @return
     */
    private int calculateFactor(TreeNode<E> node) {
        if (node == null) {
            return 0;
        }
        int factor = getHeight(node.left) - getHeight(node.right);
        return factor;
    }


    private int getHeight(TreeNode<E> node) {
        if (Objects.isNull(node)) {
            return 0;
        }
        return node.height;
    }

    /**
     * 获取最小值
     *
     * @param treeNode
     * @return
     */
    private TreeNode<E> getMin(TreeNode<E> treeNode) {

        if (Objects.isNull(treeNode.left)) {
            return treeNode;
        }

        return getMin(treeNode.left);

    }

    @Override
    public boolean remove(E element) {
        remove(root, element);
        return true;
    }

    public TreeNode<E> remove(TreeNode<E> node, E element) {

        // 1、设置退出 条件
        if (Objects.isNull(node)) {
            return null;
        }

        TreeNode<E> refNode = null;

        // 2、进行二叉树 删除操作
        // 小于，往左边查找
        if (element.compareTo(node.element) < 0) {
            refNode = remove(node.left, element);
        }
        // 大于 ，往树右边查找
        else if (element.compareTo(node.element) > 0) {
            refNode = remove(node.right, element);
        }
        // 等于，进行删除操作
        else {

            // 左节点为 null
            if (Objects.isNull(node.left)) {
                TreeNode<E> right = node.right;
                node.right = null;
                refNode = right;
            }
            // 右节点为 null
            else if (Objects.isNull(node.right)) {
                TreeNode<E> left = node.left;
                node.left = null;
                refNode = left;
            }
            // 左右都不能为null
            else {

                TreeNode<E> min = getMin(node.right);
                min.right = remove(node, element);
                min.left = node.left;
                // gc
                node.left = node.right = null;
                refNode = min;
            }
        }

        // 3、 对删除后的节点，进行平衡操作
        TreeNode<E> balance = balance(refNode);

        return balance;
    }


    @Override
    public E get(int index) {
        return null;
    }

    @Override
    public E getLast() {
        return null;
    }

    /**
     * 中序遍历 , 可以实现排序
     *
     * @return
     */
    public java.util.List<E> srot() {
        return middleTraversal(new ArrayList<E>(), root);
    }

    public java.util.List<E> middleTraversal(java.util.List<E> list, TreeNode<E> node) {

        if (node == null) {
            return null;
        }

        middleTraversal(list, node.left);
        list.add(node.element);
        middleTraversal(list, node.right);

        return list;
    }


    public static class TreeNode<E extends Comparable<E>> {
        private TreeNode<E> left, right;
        private E element;
        /**
         * 高度
         */
        private int height = 1;

        public TreeNode(E element) {
            this.element = element;
        }
    }

    @Override
    public String toString() {


        return null;


    }

    public static void main(String[] args) {

        AvlTree<Integer> avlTree = new AvlTree<>();

//        avlTree.add(41);
//        avlTree.add(22);
//        avlTree.add(60);
//        avlTree.add(15);
//        avlTree.add(33);
//        avlTree.add(59);
//        avlTree.add(63);
//        avlTree.add(13);
//        avlTree.add(37);

        avlTree.add(41);
        avlTree.add(40);
        avlTree.add(39);


    }
}
